package org.example.yolodemo.util;

import ai.onnxruntime.*;
import com.alibaba.fastjson.JSONObject;
import org.opencv.core.*;
import org.opencv.imgcodecs.Imgcodecs;
import org.opencv.imgproc.Imgproc;

import java.io.File;
import java.io.IOException;
import java.net.URL;
import java.net.URLDecoder;
import java.nio.FloatBuffer;
import java.text.DecimalFormat;
import java.util.*;
import java.util.List;

/**
 * ONNX模型加载与目标检测工具类
 */
public class OnnxLoadUtil {
    public static OrtEnvironment env; // ONNX运行时环境
    public static OrtSession session; // ONNX模型会话
    public static JSONObject names; // 存储类别名称的JSON对象
    public static long count; // 模型每次处理的图片数量（通常为1）
    public static long channels; // 模型输入通道数（通常为3）
    public static long netHeight; // 模型输入高度
    public static long netWidth; // 模型输入宽度
    public static float srcw; // 原始图片宽度
    public static float srch; // 原始图片高度
    public static float confThreshold = 0.25f; // 置信度阈值，用于过滤低置信度检测框
    public static float nmsThreshold = 0.5f; // 非极大值抑制阈值
    static Mat src; // 存储原始输入图片的Mat对象

    /**
     * 加载ONNX模型及相关配置
     * @param path ONNX模型文件路径
     */
    public static void load(String path) {
        try {
            // 加载ONNX模型
            env = OrtEnvironment.getEnvironment(); // 获取ONNX运行时环境
            session = env.createSession(path, new OrtSession.SessionOptions()); // 创建模型会话
            OnnxModelMetadata metadata = session.getMetadata(); // 获取模型元数据
            String nameClass = metadata.getCustomMetadata().get("names"); // 从元数据中获取类别名称字符串
            names = JSONObject.parseObject(nameClass.replace("\"", "\"\"")); // 解析类别名称字符串为JSON对象

            //从 session 对象中获取输入信息，并将其存储在一个键值对映射（Map）中。
            Map<String, NodeInfo> infoMap = session.getInputInfo(); // 获取输入信息映射
            TensorInfo nodeInfo = (TensorInfo) infoMap.get("images").getInfo(); // 获取输入节点“images”的信息
            count = nodeInfo.getShape()[0]; // 获取模型每次处理的图片数量（通常为1）
            channels = nodeInfo.getShape()[1]; // 获取模型输入通道数（通常为3）
            netHeight = nodeInfo.getShape()[2]; // 获取模型输入高度
            netWidth = nodeInfo.getShape()[3]; // 获取模型输入宽度

            System.loadLibrary(Core.NATIVE_LIBRARY_NAME); // 加载OpenCV本地动态库
        } catch (Exception e) {
            e.printStackTrace();
            System.exit(0); // 加载失败则退出程序
        }
    }

    /**
     * 对指定路径的图片进行预测
     * @param path 图片文件路径
     * @return 包含检测结果（框信息和类别名称）的Map对象
     * @throws Exception 可能抛出的异常
     */
    public static Map<Object, Object> predict(String path) throws Exception {
        src = Imgcodecs.imread(path); // 使用OpenCV读取图片到Mat对象
        return predictor(); // 调用具体的预测处理方法
    }

    /**
     * 将OpenCV的Mat对象转换为ONNX张量
     * @param dst 输入的Mat对象（已调整大小）
     * @return 转换后的ONNX张量
     */
    public static OnnxTensor transferTensor(Mat dst) {
        Imgproc.cvtColor(dst, dst, Imgproc.COLOR_BGR2RGB); // 将颜色空间从BGR转换为RGB
        dst.convertTo(dst, CvType.CV_32FC1, 1. / 255); // 转换图像类型为32位浮点型并归一化（0-1）
        // 创建存储图像数据的数组，大小为通道数×宽度×高度
        float[] whc = new float[Long.valueOf(channels).intValue() * Long.valueOf(netWidth).intValue() * Long.valueOf(netHeight).intValue()];
        dst.get(0, 0, whc); // 从Mat对象中获取图像数据到数组
        float[] chw = whc2cwh(whc); // 转换数组结构（从宽高通道到通道宽高）
        OnnxTensor tensor = null;
        try {
            // 创建ONNX张量，指定数据、环境和形状（[count, channels, netWidth, netHeight]）
            tensor = OnnxTensor.createTensor(env, FloatBuffer.wrap(chw), new long[]{count, channels, netWidth, netHeight});
        } catch (Exception e) {
            e.printStackTrace();
            System.exit(0); // 转换失败则退出程序
        }
        return tensor;
    }

    /**
     * 将宽高通道格式的数组转换为通道宽高格式
     * @param src 输入的宽高通道格式数组
     * @return 转换后的通道宽高格式数组
     */
    public static float[] whc2cwh(float[] src) {
        float[] chw = new float[src.length];
        int j = 0;
        for (int ch = 0; ch < 3; ++ch) { // 遍历每个通道
            for (int i = ch; i < src.length; i += 3) { // 按通道顺序重组数据
                chw[j] = src[i];
                j++;
            }
        }
        return chw;
    }

    /**
     * 执行具体的预测和后处理（非极大值抑制等）
     * @return 包含检测结果（框信息和类别名称）的Map对象
     * @throws Exception 可能抛出的异常
     */
    public static Map<Object, Object> predictor() throws Exception {
        srcw = src.width(); // 获取原始图片宽度
        srch = src.height(); // 获取原始图片高度
        float scaleW = srcw / netWidth; // 计算宽度缩放比例
        float scaleH = srch / netHeight; // 计算高度缩放比例
        Mat dst = new Mat();
        Imgproc.resize(src, dst, new Size(netWidth, netHeight)); // 调整图片大小以适应模型输入
        OnnxTensor tensor = transferTensor(dst); // 转换为ONNX张量
        OrtSession.Result result = session.run(Collections.singletonMap("images", tensor)); // 运行模型推理
        System.out.println("res Data: " + result.get(0)); // 打印推理结果（调试用）
        OnnxTensor res = (OnnxTensor) result.get(0); // 获取推理结果张量

        float[][][] dataRes = (float[][][]) res.getValue(); // 获取推理结果值
        float[][] data = dataRes[0]; // 取第一个元素（因为count=1）

        // 创建原始数据数组，每行存储一个检测框的信息（x, y, w, h, 置信值, 类别编号）
        float rawData[][] = new float[data[0].length][6];
        System.out.println("length="+(data.length - 1)); // 打印数据长度（调试用）
        for (int i = 0; i < 4; i++) { // 转置xywh部分数据（将列转成行）
            for (int j = 0; j < data[0].length; j++) {
                rawData[j][i] = data[i][j];
            }
        }

        // 寻找每个检测框置信度最高的类别
        for (int i = 0; i < data[0].length; i++) {
            for (int j = 4; j < data.length; j++) {
                if (rawData[i][4] < data[j][i]) {
                    rawData[i][4] = data[j][i]; // 更新置信值
                    rawData[i][5] = j - 4; // 更新类别编号
                }
            }
        }

        List<ArrayList<Float>> boxes = new LinkedList<ArrayList<Float>>(); // 存储检测框的列表
        ArrayList<Float> box = null;
        // 过滤低置信度检测框，并将xywh转换为xyxy（左上角和右下角坐标）
        for (float[] d : rawData) {
            if (d[4] > confThreshold) { // 置信度过滤
                // xywh（中心点x, y，宽, 高）转xyxy（左上角x1, y1，右下角x2, y2）
                d[0] = d[0] - d[2] / 2;
                d[1] = d[1] - d[3] / 2;
                d[2] = d[0] + d[2];
                d[3] = d[1] + d[3];
                box = new ArrayList<Float>();
                for (float num : d) {
                    box.add(num);
                }
                // 按置信值插入排序，确保高置信度框在前
                if (boxes.size() == 0) {
                    boxes.add(box);
                } else {
                    int i;
                    for (i = 0; i < boxes.size(); i++) {
                        if (box.get(4) > boxes.get(i).get(4)) {
                            boxes.add(i, box);
                            break;
                        }
                    }
                    if (i == boxes.size()) { // 插入到列表末尾
                        boxes.add(box);
                    }
                }
            }
        }

        int[] indexs = new int[boxes.size()]; // 用于非极大值抑制的标记数组，1表示保留，0表示删除
        Arrays.fill(indexs, 1); // 初始化都为保留
        for (int cur = 0; cur < boxes.size(); cur++) {
            if (indexs[cur] == 0) {
                continue;
            }
            ArrayList<Float> curMaxConf = boxes.get(cur); // 当前置信度最高的检测框
            for (int i = cur + 1; i < boxes.size(); i++) {
                if (indexs[i] == 0) {
                    continue;
                }
                float classIndex = boxes.get(i).get(5);
                // 仅当两个检测框属于同一类别时进行非极大值抑制
                if (classIndex == curMaxConf.get(5)) {
                    float x1 = curMaxConf.get(0);
                    float y1 = curMaxConf.get(1);
                    float x2 = curMaxConf.get(2);
                    float y2 = curMaxConf.get(3);
                    float x3 = boxes.get(i).get(0);
                    float y3 = boxes.get(i).get(1);
                    float x4 = boxes.get(i).get(2);
                    float y4 = boxes.get(i).get(3);
                    // 排除不相交的情况
                    if (x1 > x4 || x2 < x3 || y1 > y4 || y2 < y3) {
                        continue;
                    }
                    // 计算交集宽度
                    float intersectionWidth = Math.max(x1, x3) - Math.min(x2, x4);
                    // 计算交集高度
                    float intersectionHeight = Math.max(y1, y3) - Math.min(y2, y4);
                    float intersectionArea = Math.max(0, intersectionWidth * intersectionHeight); // 交集面积
                    float unionArea = (x2 - x1) * (y2 - y1) + (x4 - x3) * (y4 - y3) - intersectionArea; // 并集面积
                    float iou = intersectionArea / unionArea; // 计算交并比（IoU）
                    // 根据IoU和阈值标记是否保留
                    indexs[i] = iou > nmsThreshold? 0 : 1;
                }
            }
        }

        List<ArrayList<Float>> resBoxes = new LinkedList<ArrayList<Float>>(); // 存储最终保留的检测框
        for (int index = 0; index < indexs.length; index++) {
            if (indexs[index] == 1) {
                resBoxes.add(boxes.get(index));
            }
        }
        boxes = resBoxes; // 更新检测框列表
        // 还原检测框坐标到原始图片尺寸
        for (ArrayList<Float> box1 : boxes) {
            box1.set(0, box1.get(0) * scaleW);
            box1.set(1, box1.get(1) * scaleH);
            box1.set(2, box1.get(2) * scaleW);
            box1.set(3, box1.get(3) * scaleH);
        }
        Map<Object, Object> map = new HashMap<Object, Object>();
        map.put("boxes", boxes); // 将检测框存入Map
        map.put("classNames", names); // 将类别名称存入Map
        return map;
    }

    /**
     * 显示检测结果（画框和类别信息）并保存结果图片
     * @param map 包含检测框（boxes）和类别名称（classNames）的Map对象
     */
    public static void showDetect(Map<Object, Object> map) {
        List<ArrayList<Float>> boxes = (List<ArrayList<Float>>) map.get("boxes"); // 获取检测框列表
        JSONObject names = (JSONObject) map.get("classNames"); // 获取类别名称JSON对象
        Imgproc.resize(src, src, new Size(srcw, srch)); // 将图片恢复为原始大小
        // 遍历每个检测框并绘制矩形和文字
        for (ArrayList<Float> box : boxes) {
            float x1 = box.get(0);
            float y1 = box.get(1);
            float x2 = box.get(2);
            float y2 = box.get(3);
            float config = box.get(4);
            String className = (String) names.get((int) box.get(5).intValue()); // 获取类别名称
            Point point1 = new Point(x1, y1);
            Point point2 = new Point(x2, y2);
            Imgproc.rectangle(src, point1, point2, new Scalar(0, 0, 255), 2); // 绘制红色矩形框
            String conf = new DecimalFormat("#.###").format(config); // 格式化置信度为三位小数
            // 在框上方绘制类别名称和置信度
            Imgproc.putText(src, className + " " + conf, new Point(x1, y1 - 5), 0, 0.8, new Scalar(255, 0, 0), 2);
        }

        // 保存结果图片到指定路径
        boolean success = Imgcodecs.imwrite("C:\\Users\\HP\\IdeaProjects\\YoloDemo\\results\\result.jpg", src);
        if (!success) {
            System.err.println("保存图片失败: " + "C:\\Users\\HP\\IdeaProjects\\YoloDemo\\results\\result.jpg");
        }
    }

    /**
     * 获取模型文件的实际路径（处理资源路径和编码）
     * @param resourcePath 模型资源路径（类路径下的路径）
     * @return 模型文件的实际路径
     * @throws IOException 路径获取或文件检查时可能抛出的异常
     */
    public static String getModelPath(String resourcePath) throws IOException {
        URL resourceUrl = ClassLoader.getSystemResource(resourcePath); // 获取资源URL  models/best.onnx
        if (resourceUrl == null) {
            throw new IOException("资源未找到: " + resourcePath);
        }
        String path = URLDecoder.decode(resourceUrl.getPath(), "UTF-8"); // 解码URL路径
        ///D:/IdeaProjects/yoloOpen/target/classes/models/best.onnx
        System.out.println("////////---->"+path);
        // 处理Windows系统下路径前导斜杠问题
        if (System.getProperty("os.name").toLowerCase().contains("win") && path.startsWith("/")) {
            path = path.substring(1);  //因为路径的前边有一个/,所以要去掉它
            System.out.println("------------->"+path);
        }
        //D:/IdeaProjects/yoloOpen/target/classes/models/best.onnx
        File file = new File(path);
        if (!file.exists()) {
            throw new IOException("文件不存在: " + path);
        }
        return path;
    }

    /**
     * 主方法（测试用）
     * @param args 命令行参数（未使用）
     * @throws Exception 可能抛出的异常
     */
    public static void main(String[] args) throws Exception {
        String modelPath = getModelPath("models/best.onnx"); // 获取模型文件路径
        String path = "C:\\Users\\HP\\IdeaProjects\\YoloDemo\\Image\\second.jpg"; // 测试图片路径
        OnnxLoadUtil.load(modelPath); // 加载ONNX模型
        Map<Object, Object> map = OnnxLoadUtil.predict(path); // 对测试图片进行预测
        showDetect(map); // 显示检测结果并保存图片
    }
}